"""
基于角色的LLM处理系统
参考 /root/code/drwork/aimeeting/backend/Role/RoleList.py 实现
"""
import os
import time
from openai import OpenAI
from typing import Optional


class KimiClient:
    """Kimi API客户端"""
    def __init__(self, api_key: str, model='kimi-k2-turbo-preview', try_times=5):
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://api.moonshot.cn/v1"
        )
        self.model = model
        self.try_times = try_times

    def request(self, text: str, role_prompt: str) -> str:
        """向Kimi API发送请求"""
        for times in range(self.try_times):
            try:
                completions = self.client.chat.completions.create(
                    model=self.model,
                    messages=[
                        {'role': 'system', 'content': role_prompt},
                        {'role': 'user', 'content': text}
                    ]
                )
                return completions.choices[0].message.content
            except Exception as e:
                print(f'WARNING: 连接Kimi API失败: {e}')
                time.sleep((times + 1) * 5)
            if times == self.try_times - 1:
                raise Exception(f'尝试连接KIMI API {self.try_times} 次均失败，请检查网络连接或API Key。')


class Role:
    """角色基类"""
    def __init__(self, role_file: str, model: KimiClient) -> None:
        self.model = model
        self.role_prompt = self._read_role_prompt(role_file)

    def _read_role_prompt(self, role_file: str) -> str:
        """读取角色提示词"""
        current_dir = os.path.dirname(os.path.abspath(__file__))
        with open(f'{current_dir}/prompts/{role_file}', 'r', encoding='utf-8') as f:
            return f.read()

    def request(self, text: str) -> str:
        """使用角色处理文本"""
        return self.model.request(text, self.role_prompt)


class MeetingSecretary(Role):
    """会议秘书角色"""
    def __init__(self, model: KimiClient, template: Optional[str] = None) -> None:
        super().__init__('MeetingSecretary.md', model)
        self.raw_prompt = self.role_prompt
        self.set_template(template if template is not None else self._get_default_template())

    def set_template(self, template: str) -> None:
        """设置模板并注入到提示词中"""
        self.role_prompt = self.raw_prompt.replace('${template}', template)

    def _get_default_template(self) -> str:
        """获取默认模板"""
        current_dir = os.path.dirname(os.path.abspath(__file__))
        with open(f'{current_dir}/prompts/DefaultMinutesTemplate.md', 'r', encoding='utf-8') as f:
            return f.read()


class SummaryWriter(Role):
    """总结撰写角色"""
    def __init__(self, model: KimiClient) -> None:
        super().__init__('SummaryWriter.md', model)


class MeetingMinutesEditor(Role):
    """会议纪要编辑角色"""
    def __init__(self, model: KimiClient, template: Optional[str] = None) -> None:
        super().__init__('MeetingMinutesEditor.md', model)
        self.raw_prompt = self.role_prompt
        self.set_template(template if template is not None else self._get_default_template())

    def set_template(self, template: str) -> None:
        """设置模板并注入到提示词中"""
        self.role_prompt = self.raw_prompt.replace('${template}', template)

    def _get_default_template(self) -> str:
        """获取默认模板"""
        current_dir = os.path.dirname(os.path.abspath(__file__))
        with open(f'{current_dir}/prompts/DefaultMinutesTemplate.md', 'r', encoding='utf-8') as f:
            return f.read()


class RoleBasedProcessor:
    """基于角色的多阶段处理器"""
    def __init__(self, api_key: str, template: Optional[str] = None):
        self.model = KimiClient(api_key)
        self.template = template
        self.summary_writer = SummaryWriter(self.model)
        self.meeting_secretary = MeetingSecretary(self.model, template)
        self.meeting_minutes_editor = MeetingMinutesEditor(self.model, template)

    def process_transcript_to_minutes(self, transcript_text: str, progress_callback=None) -> str:
        """
        将转录文本转换为会议纪要
        采用三阶段处理：
        1. 总结生成 - 将长文本分段总结
        2. 纪要生成 - 基于总结生成结构化纪要
        3. 纪要编辑 - 完善和优化纪要内容
        """
        if progress_callback:
            progress_callback(10)

        # 阶段1: 生成总结（分段处理）
        print("开始生成会议总结...")
        summary = self._generate_summary(transcript_text, progress_callback)

        if progress_callback:
            progress_callback(50)

        # 阶段2: 生成会议纪要
        print("基于总结生成会议纪要...")
        minutes_draft = self.meeting_secretary.request(summary)

        if progress_callback:
            progress_callback(80)

        # 阶段3: 编辑和完善会议纪要
        print("编辑和完善会议纪要...")
        final_minutes = self.meeting_minutes_editor.request(
            f"## 原始会议内容 \n{summary} \n\n## 会议纪要草稿 \n{minutes_draft}"
        )

        if progress_callback:
            progress_callback(100)

        return final_minutes

    def _generate_summary(self, text: str, progress_callback=None) -> str:
        """将长文本分段总结"""
        clip_length = 6000  # 每段长度
        clip_abandoned = 1000  # 最短片段长度

        # 修复分段逻辑：如果文本太短，至少作为1段处理
        if len(text) <= clip_abandoned:
            clip_count = 1
        else:
            clip_count = len(text) // clip_length + 1 if len(text) % clip_length >= clip_abandoned else len(text) // clip_length

        clips = []
        for i in range(clip_count):
            start = clip_length * i
            end = min(clip_length * (i + 1), len(text))
            clips.append(text[start:end])

        print(f'文本共 {len(text)} 字符，将分为 {clip_count} 段进行处理。')

        summary = ''
        for i, clip in enumerate(clips):
            print(f'正在处理第 {i + 1}/{clip_count} 段...')
            segment_summary = self.summary_writer.request(clip)
            summary += f'## 第{i + 1}段内容总结: \n{segment_summary}\n\n'

            if progress_callback:
                # 阶段1占总进度的40%
                segment_progress = 10 + (40 * (i + 1) / clip_count)
                progress_callback(int(segment_progress))

        return summary